/*
 * Copyright(C) 2022. Huawei Technologies Co.,Ltd. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <iostream>
#include <vector>
#include <algorithm>
#include <map>
#include "opencv2/opencv.hpp"
#include "MxBase/Log/Log.h"
#include "MxBase/Maths/FastMath.h"
#include "MxBase/MxBase.h"
#include "MxBase/postprocess/include/ObjectPostProcessors/Yolov3PostProcess.h" // sdk inside YoloV3PostProcess fuc


#define USE_DVPP // if use opencv, this options should disable
//#define USE_200DK // only enable when using 200DK
using namespace MxBase;
using namespace std;
namespace {
    const uint32_t YUV_BYTE_NU = 3;
    const uint32_t YUV_BYTE_DE = 2;
    const uint32_t VPC_H_ALIGN = 2;
    const uint32_t YOLOV3_RESIZE = 416;
}
struct V2Param {
    uint32_t deviceId;
    std::string labelPath;
    std::string configPath;
    std::string modelPath;
};

void InitV2Param(V2Param &v2Param)
{
    v2Param.deviceId = 0;
    v2Param.labelPath = "./model/yolov3.names";
    v2Param.configPath = "./model/yolov3_tf_bs1_fp16.cfg";
    v2Param.modelPath = "./model/yolov3_tf_bs1_fp16.om";
};

APP_ERROR GetImage(const std::string &imgPath, std::shared_ptr<uint8_t>& dataPtr, uint32_t& dataSize)
{
    // Get image data to memory, this method can be substituted or designed by yourself!
    std::ifstream file;
    file.open(imgPath.c_str(), std::ios::binary);
    if (!file)
    {
        std::cout << "Invalid file." << std::endl;
        return APP_ERR_COMM_INVALID_PARAM;
    }
    std::stringstream buffer;
    buffer << file.rdbuf();
    std::string content = buffer.str();

    char* p = (char*) malloc(content.size());
    memcpy(p, content.data(), content.size());
    auto deleter = [](void *p) -> void {
        free(p);
        p = nullptr;
    };

    dataPtr.reset(static_cast<uint8_t*>((void*)(p)), deleter);
    dataSize = content.size();

    file.close();
    return APP_ERR_OK;
}

void YoloV3PostProcess(std::string& yoloV3ConfigPath, std::string& yoloV3LablePath,
                        std::vector<MxBase::Tensor>& yoloV3Outputs, std::vector<MxBase::Rect>& cropConfigVec)
{
    /// This should made by user! This func only show used with sdk`s yolov3 so lib.

    std::cout << "******YoloV3PostProcess******" << std::endl;
    // make yoloV3config map
    std::map<std::string, std::string> postConfig;

    postConfig.insert(pair<std::string, std::string>("postProcessConfigPath", yoloV3ConfigPath));
    postConfig.insert(pair<std::string, std::string>("lablePath", yoloV3LablePath));

    // init postProcess
    Yolov3PostProcess yolov3PostProcess;
    yolov3PostProcess.Init(postConfig);

    // make postProcess inputs
    vector<TensorBase> tensors;
    for (size_t i = 0; i < yoloV3Outputs.size(); i++)
    {
        MemoryData memoryData(yoloV3Outputs[i].GetData(), yoloV3Outputs[i].GetByteSize());
        TensorBase tensorBase(memoryData, true, yoloV3Outputs[i].GetShape(), TENSOR_DTYPE_INT32);
        tensors.push_back(tensorBase);
    }
    vector<vector<ObjectInfo>> objectInfos;

    auto shape = yoloV3Outputs[0].GetShape();
    MxBase::ResizedImageInfo imgInfo;
    imgInfo.widthOriginal = shape[1];
    imgInfo.heightOriginal = shape[0] * YUV_BYTE_DE / YUV_BYTE_NU;
    imgInfo.widthResize = YOLOV3_RESIZE;
    imgInfo.heightResize = YOLOV3_RESIZE;
    imgInfo.resizeType = MxBase::RESIZER_STRETCHING;
    std::vector<MxBase::ResizedImageInfo> imageInfoVec = {};
    imageInfoVec.push_back(imgInfo);
    
    // do postProcess
    yolov3PostProcess.Process(tensors, objectInfos, imageInfoVec);

    // print result
    std::cout << "Size of objectInfos is " << objectInfos.size() << std::endl;
    for (size_t i = 0; i < objectInfos.size(); i++)
    {
        std::cout << "objectInfo-" << i << " ,Size:"<< objectInfos[i].size() << std::endl;
        for (size_t j = 0; j < objectInfos[i].size(); j++)
        {
            std::cout << std::endl << "*****objectInfo-" << i << ":" << j << std::endl;
            std::cout << "x0 is " << objectInfos[i][j].x0 << std::endl;
            std::cout << "y0 is " << objectInfos[i][j].y0 << std::endl;
            std::cout << "x1 is " << objectInfos[i][j].x1 << std::endl;
            std::cout << "y1 is " << objectInfos[i][j].y1 << std::endl;
            std::cout << "confidence is " << objectInfos[i][j].confidence << std::endl;
            std::cout << "classId is " << objectInfos[i][j].classId << std::endl;
            std::cout << "className is " << objectInfos[i][j].className << std::endl;           
        }
    }
    
    // do image crop
    cropConfigVec.resize(objectInfos[0].size());
    for (size_t i = 0; i < objectInfos[0].size(); i++)
    {
        cropConfigVec[i].x0 = objectInfos[0][i].x0;
        cropConfigVec[i].y0 = objectInfos[0][i].y0;
        cropConfigVec[i].x1 = objectInfos[0][i].x1;
        cropConfigVec[i].y1 = objectInfos[0][i].y1;
    }
    
    std::cout << "******YoloV3PostProcess end******" << std::endl;
};

void generalPostProcess(std::vector<MxBase::Tensor>& inferOutputs)
{
    // Modify this func to finish your postprocess, need add return
    size_t count = 0; // define result count or got by .Size()
    if (inferOutputs.empty())
    {
        std::cout << "result Infer failed with empty output..." << std::endl;
        return;
    }
    
    // convert tensor to cpp data. This sample convert first result tensor to float32 array
    float *castData = static_cast<float *>(inferOutputs[0].GetData());
    
    std::vector<float> result_0;
    for (size_t i = 0; i < count; i++)
    {
        result_0.push_back(castData[i]);
    }

    /// do compute and add return

    return;
}

int main(int argc, char* argv[])
{
    if (argc <= 1) 
    {
        LogWarn << "Please input image path, such as './cppv2_sample test.jpg'.";
        return APP_ERR_OK;
    }
    APP_ERROR ret;
    
    // *****1初始化模型推理
    V2Param v2Param;
    InitV2Param(v2Param);
    int32_t deviceId = v2Param.deviceId;
    std::string modelPath = v2Param.modelPath;

    // global init
    ret = MxInit();
    if (ret != APP_ERR_OK) 
    {
        LogError << "MxInit failed, ret=" << ret << ".";
    }

    // imageProcess init
    MxBase::ImageProcessor imageProcessor(deviceId); 

    // model init
    MxBase::Model yoloV3(modelPath, deviceId);

    // *****2读取图片
    std::string imgPath = argv[1]; 
#ifdef USE_DVPP
    MxBase::Image decodedImage;
#ifdef USE_200DK
    // 200DK only support this method!
    std::shared_ptr<uint8_t> dataPtr;
    uint32_t dataSize;
    Image imageData;
    ret = GetImage(imgPath, dataPtr, dataSize);
    if (ret != APP_ERR_OK) 
    {
        LogError << "Getimage failed, ret=" << ret;
        return ret;
    }
    ret = imageProcessor.Decode(dataPtr, dataSize, decodedImage, ImageFormat::YUV_SP_420);
#else
    ret = imageProcessor.Decode(imgPath, decodedImage, ImageFormat::YUV_SP_420);
#endif
    if (ret != APP_ERR_OK) 
    {
        LogError << "Decode failed, ret=" << ret;
        return ret;
    }
    LogInfo << "decodedImage.width is" << decodedImage.GetOriginalSize().width;
    LogInfo << "decodedImage.height is" << decodedImage.GetOriginalSize().height;
    LogInfo << "decodedImage.format is" << (int)(decodedImage.GetFormat());

    // *****3缩放图片
    MxBase::Image resizeImage;
    // set size param
    Size resizeConfig(YOLOV3_RESIZE, YOLOV3_RESIZE);

    ret = imageProcessor.Resize(decodedImage, resizeConfig, resizeImage, Interpolation::HUAWEI_HIGH_ORDER_FILTER);
    if (ret != APP_ERR_OK) 
    {
        LogError << "Resize failed, ret=" << ret;
        return ret;
    }

    // save resize image
    std::string path = "./resized_yolov3_416.jpg";
    ret = imageProcessor.Encode(resizeImage, path);
    if (ret != APP_ERR_OK) 
    {
        LogError << "Encode failed, ret=" << ret;
        return ret;
    }

    // *****4模型推理
    Tensor tensorImg = resizeImage.ConvertToTensor();
    ret = tensorImg.ToDevice(deviceId);
    if (ret != APP_ERR_OK) 
    {
        LogError << "ToDevice failed, ret=" << ret;
        return ret;
    }

    // make infer input
    std::vector<Tensor> yoloV3Inputs = {tensorImg};
    // do infer
    std::vector<Tensor> yoloV3Outputs = yoloV3.Infer(yoloV3Inputs);
    std::cout << "yoloV3Outputs len=" << yoloV3Outputs.size() << std::endl;

    // !move result to host!
    for (auto output : yoloV3Outputs)
    {
        output.ToHost();
    }

#else // use opencv 替代L159-L228
    // *****2opencv读取
    cv::Mat srcImg = imread(imgPath, cv::IMREAD_COLOR);
    if (srcImg.data == nullptr)
    {
        LogError << "opencv readImage failed";
    }
    cv::cvtColor(srcImg, srcImg, cv::COLOR_BGR2RGB);    // 需要和模型输入相同，此处对应RGB模型（本样例中使用了AIPP的YOLOV3模型为YUV输入，因此直接使用无推理结果）

    // *****3opecnv缩放
    cv::Mat dstImg;
    cv::resize(srcImg, dstImg, cv::Size(YOLOV3_RESIZE, YOLOV3_RESIZE));

/*
    使用opencv进行crop示例
    static cv::Rect rectOfImg(x1, y1, x2, y2)
*/

    // ******4模型推理
    // cv::mat to Tensor
    std::vector<uint32_t> shape;
    shape.push_back(1); //batchSize
    shape.push_back(dstImg.rows);
    shape.push_back(dstImg.cols);
    shape.push_back(dstImg.channels());

    MxBase::TensorDType tensorDataType = MxBase::TensorDType::UINT8;
    MxBase::Tensor tensorInput(dstImg.data, shape, tensorDataType); 

    tensorInput.ToDevice(v2Param.deviceId); // !!!!!重要，使用外部数据作为tensor时务必使用to_device进行转移，缺失该步骤会导致输出结果异常，RC3以上版本已修复
    std::vector<MxBase::Tensor> yoloV3Inputs = {};
    yoloV3Inputs.push_back(tensorInput);


    // do infer
    std::vector<Tensor> yoloV3Outputs = yoloV3.Infer(yoloV3Iutputs);
    std::cout << "yoloV3Outputs len=" << yoloV3Outputs.size() << std::endl;

    // !move result to host!
    for (auto output : yoloV3Outputs)
    {
        output.ToHost();
    }
#endif

    // *****5模型后处理
    std::vector<MxBase::Rect> cropConfigVec;
    YoloV3PostProcess(v2Param.configPath, v2Param.labelPath, yoloV3Outputs, cropConfigVec);

    return APP_ERR_OK;
};
